# -*- coding: utf-8 -*-
"""
数据采集器
负责从元数据库采集数据
"""

import logging
from typing import Dict, List, Any, Optional
import pymysql
from pymysql.cursors import DictCursor
from dbutils.pooled_db import PooledDB

logger = logging.getLogger(__name__)


class DataCollector:
    """数据采集器"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.db_pool = self._create_db_pool()
        logger.info("数据采集器初始化完成")
    
    def _create_db_pool(self) -> PooledDB:
        """创建数据库连接池"""
        mysql_config = self.config.get('mysql', {})
        return PooledDB(
            creator=pymysql,
            maxconnections=mysql_config.get('pool_size', 5),
            host=mysql_config.get('host', 'localhost'),
            port=mysql_config.get('port', 3306),
            user=mysql_config.get('user'),
            password=mysql_config.get('password'),
            database=mysql_config.get('database', 'hive_metastore'),
            charset='utf8mb4',
            cursorclass=DictCursor
        )
    
    def _get_connection(self):
        """获取数据库连接"""
        return self.db_pool.connection()
    
    def get_all_databases(self) -> List[str]:
        """获取所有数据库"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cursor:
                cursor.execute("""
                    SELECT NAME as db_name 
                    FROM DBS 
                    WHERE NAME NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')
                """)
                return [row['db_name'] for row in cursor.fetchall()]
        finally:
            conn.close()
    
    def get_all_tables(self) -> List[Dict[str, Any]]:
        """获取所有表信息"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cursor:
                cursor.execute("""
                    SELECT 
                        d.NAME as db_name,
                        t.TBL_NAME as table_name,
                        t.TBL_ID as table_id,
                        t.TBL_TYPE as table_type,
                        t.CREATE_TIME as create_time,
                        t.OWNER as owner
                    FROM TBLS t
                    JOIN DBS d ON t.DB_ID = d.DB_ID
                    WHERE d.NAME NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')
                """)
                return cursor.fetchall()
        finally:
            conn.close()
    
    def get_table_details(self, table_id: int) -> Dict[str, Any]:
        """获取表详细信息"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cursor:
                # 获取分区数
                cursor.execute("""
                    SELECT COUNT(*) as partition_count
                    FROM PARTITIONS
                    WHERE TBL_ID = %s
                """, (table_id,))
                partition_result = cursor.fetchone()
                partition_count = partition_result['partition_count'] if partition_result else 0
                
                # 获取表参数
                cursor.execute("""
                    SELECT PARAM_KEY, PARAM_VALUE
                    FROM TABLE_PARAMS
                    WHERE TBL_ID = %s AND PARAM_KEY IN ('numRows', 'totalSize', 'numFiles')
                """, (table_id,))
                params = {p['PARAM_KEY']: p['PARAM_VALUE'] for p in cursor.fetchall()}
                
                return {
                    'partition_count': partition_count,
                    'row_count': int(params.get('numRows', 0)),
                    'total_size': int(params.get('totalSize', 0)),
                    'num_files': int(params.get('numFiles', 0))
                }
        finally:
            conn.close()
    
    def get_database_table_count(self, db_name: str) -> int:
        """获取数据库下的表数量"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cursor:
                cursor.execute("""
                    SELECT COUNT(*) as cnt
                    FROM TBLS t
                    JOIN DBS d ON t.DB_ID = d.DB_ID
                    WHERE d.NAME = %s
                """, (db_name,))
                result = cursor.fetchone()
                return result['cnt'] if result else 0
        finally:
            conn.close()
    
    def test_connection(self) -> bool:
        """测试数据库连接"""
        try:
            conn = self._get_connection()
            with conn.cursor() as cursor:
                cursor.execute("SELECT 1")
                cursor.fetchone()
            conn.close()
            return True
        except Exception as e:
            logger.error(f"数据库连接测试失败: {e}")
            return False

